package com.boat.ai.adaptor;

import com.boat.ai.FileTypeUtil;
import com.boat.ai.domain.entity.KwDocument;
import com.boat.ai.domain.entity.KwSegment;
import com.boat.ai.domain.enums.FileTypeEnum;
import com.boat.ai.dto.req.KwSegmentReq;
import com.boat.ai.dto.req.TextAddReq;
import com.boat.ai.dto.rsp.DocContentRsp;
import com.boat.ai.dto.rsp.KwSegmentRsp;
import com.boat.ai.dto.rsp.R;
import com.boat.ai.rag.Ingestion.strategy.ParseStrategy;
import com.boat.ai.rag.Ingestion.strategy.ParseStrategyFactory;
import com.boat.ai.rag.embedding.IngestWithQdrant;
import com.boat.ai.util.IdUtil;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.segment.TextSegment;
import jakarta.inject.Inject;
import jakarta.transaction.Transactional;
import jakarta.validation.Valid;
import jakarta.validation.constraints.NotNull;
import jakarta.ws.rs.*;
import jakarta.ws.rs.core.MediaType;
import org.eclipse.microprofile.openapi.annotations.Operation;
import org.eclipse.microprofile.openapi.annotations.tags.Tag;
import org.jboss.resteasy.reactive.RestForm;
import org.jboss.resteasy.reactive.RestQuery;
import org.jboss.resteasy.reactive.multipart.FileUpload;

import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.time.Instant;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Tag(name = "文档管理")
@Path("/document")
@Consumes(MediaType.APPLICATION_JSON)
@Produces(MediaType.APPLICATION_JSON)
public class DocumentParserAdaptor {

    @Inject
    ParseStrategyFactory parseStrategyFactory;

    @Operation(summary = "文档解析")
    @POST
    @Path("/parse")
    @Consumes(MediaType.MULTIPART_FORM_DATA)
    public R<DocContentRsp> parse(@RestQuery String title, @RestForm("file") FileUpload fileUpload) throws IOException {
        String contentType = fileUpload.contentType();
        FileTypeEnum fileType = FileTypeUtil.getFileType(contentType);
        ParseStrategy strategy = parseStrategyFactory.getStrategy(fileType);
        java.nio.file.Path path = fileUpload.uploadedFile(); // 实际是 filePath()
        InputStream is = Files.newInputStream(path);
        List<TextSegment> parse = strategy.parse(is);
        List<KwSegmentRsp> kwSegmentRspList = new ArrayList<>();
        for (TextSegment textSegment : parse) {
            String text = textSegment.text();
            String tag = textSegment.metadata()
                    .getString("tag");
            KwSegmentRsp kwSegmentRsp = new KwSegmentRsp(tag, text);
            kwSegmentRspList.add(kwSegmentRsp);
        }
        return R.ok(new DocContentRsp(title, kwSegmentRspList));
    }

    @Inject
    IngestWithQdrant ingestWithQdrant;

    @Operation(summary = "文档保存")
    @POST
    @Path("/add")
    @Transactional(rollbackOn = Exception.class)
    public R<List<KwSegmentRsp>> add(@Valid @NotNull(message = "文档不能为空") TextAddReq textAddReq) {
        KwDocument kwDocument = new KwDocument();
        long documentId = IdUtil.getId();
        kwDocument.id = documentId;
        kwDocument.knowledgeId = textAddReq.knowledgeId;
        kwDocument.title = textAddReq.title;
        kwDocument.status = textAddReq.status;
        kwDocument.createTime = Instant.now();
        kwDocument.updateTime = Instant.now();
        kwDocument.persist();
        List<KwSegment> kwSegmentList = new ArrayList<>();
        List<Document> documentList = new ArrayList<>();
        for (KwSegmentReq docSegment : textAddReq.content) {
            KwSegment kwSegment = new KwSegment();
            kwSegment.id = IdUtil.getId();
            kwSegment.knowledgeId = textAddReq.knowledgeId;
            kwSegment.documentId = documentId;
            kwSegment.createTime = Instant.now();
            kwSegment.updateTime = Instant.now();
            kwSegment.segment = docSegment.segment();
            kwSegment.tag = docSegment.tag();
            kwSegmentList.add(kwSegment);
            Map<String, String> metadataMap = new HashMap<>();
            metadataMap.put("tag", kwSegment.tag);
            metadataMap.put("knowledgeId", String.valueOf(kwSegment.knowledgeId));
            metadataMap.put("documentId", String.valueOf(kwSegment.documentId));
            Metadata metadata = Metadata.from(metadataMap);
            Document from = Document.from(kwSegment.segment, metadata);
            documentList.add(from);
        }
        KwSegment.persist(kwSegmentList);
        ingestWithQdrant.ingest(documentList);
        return R.ok();
    }


}
